/*
 *    This program is free software; you can redistribute it and/or modify
 *    it under the terms of the GNU General Public License as published by
 *    the Free Software Foundation; either version 2 of the License, or
 *    (at your option) any later version.
 *
 *    This program is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *    GNU General Public License for more details.
 *
 *    You should have received a copy of the GNU General Public License
 *    along with this program; if not, write to the Free Software
 *    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */

#include <iostream>

#include <boost/bind.hpp>
#include <boost/date_time/posix_time/posix_time_types.hpp>

#include "sgconf.h"

SGCONF::SGCONF()
    : io_service(),
      timer(io_service),
      socket(io_service),
      iv(new unsigned char[8]),
      cryptoXMLBlock(new unsigned char[4096]),
      XMLBlock(new unsigned char[4096])
{
    memset(iv.get(), 0, 8);
}

SGCONF::SGCONF(const std::string & h,
               const std::string & p,
               const std::string & l,
               const std::string & pwd)
    : io_service(),
      timer(io_service),
      socket(io_service),
      host(h),
      port(p),
      login(l),
      password(pwd),
      iv(new unsigned char[8]),
      cryptoXMLBlock(new unsigned char[4096]),
      XMLBlock(new unsigned char[4096])
{
    memset(iv.get(), 0, 8);
    BF_set_key(&key, password.length(),
               reinterpret_cast<const unsigned char *>(password.c_str()));
}

SGCONF::~SGCONF()
{
    socket.close();
}

bool SGCONF::Perform(const std::string & arg)
{
    std::string xml(arg);

    if (!Connect()) {
        return false;
    }

    if (xml.size() % 8) {
        size_t delta = 8 - xml.size() % 8;
        xml.append(delta, ' ');
    }

    // Prepare headers
    struct REQ::HEADER reqHdr;
    strncpy(reqHdr.magic, PROTO_MAGIC, sizeof(reqHdr.magic));
    reqHdr.version = htonl(2 << 8 | 0); // Major | Minor
    memset(reqHdr.login, 0, sizeof(reqHdr.login));
    strncpy(reqHdr.login, login.c_str(), sizeof(reqHdr.login));

    struct REQ::CRYPTO_HEADER reqCryptoHdr;
    memset(reqCryptoHdr.login, 0, sizeof(reqCryptoHdr.login));
    strncpy(reqCryptoHdr.login, login.c_str(), sizeof(reqCryptoHdr.login));
    if (xml.size() % 8 == 0) {
        reqCryptoHdr.dataSize = htonl(xml.size());
    } else {
        reqCryptoHdr.dataSize = htonl((xml.size() / 8 + 1) * 8);
    }

    // Crypting header
    boost::scoped_array<unsigned char> cryptoBuffer(new unsigned char[sizeof(reqCryptoHdr)]);
    BF_cbc_encrypt(reinterpret_cast<const unsigned char *>(&reqCryptoHdr), cryptoBuffer.get(), sizeof(reqCryptoHdr), &key, iv.get(), BF_ENCRYPT);

    // Crypting data
    size_t size = xml.size();
    boost::scoped_array<unsigned char> cryptoXML(new unsigned char[size]);
    BF_cbc_encrypt(reinterpret_cast<const unsigned char *>(xml.c_str()), cryptoXML.get(), size, &key, iv.get(), BF_ENCRYPT);

    try {
        // Prepare buffers
        std::vector<const_buffer> sendBuffers;
        sendBuffers.push_back(boost::asio::buffer(&reqHdr, sizeof(reqHdr)));
        sendBuffers.push_back(boost::asio::buffer(cryptoBuffer.get(), sizeof(reqCryptoHdr)));
        sendBuffers.push_back(boost::asio::buffer(cryptoXML.get(), size));

        boost::system::error_code error;
        socket.write_some(sendBuffers, error);

        if (error) {
            throw boost::system::system_error(error);
        }

        async_read(socket,
                   boost::asio::buffer(&respHeader, sizeof(respHeader)),
                   boost::bind(&SGCONF::HandleHeader,
                               this,
                               boost::asio::placeholders::error,
                               boost::asio::placeholders::bytes_transferred));
        timer.expires_from_now(boost::posix_time::seconds(5));
        timer.async_wait(boost::bind(&SGCONF::HandleHeadersTimeout, this, boost::asio::placeholders::error));

        io_service.run();
    }
    catch (std::logic_error & err) {
        std::cerr << err.what() << "\n";
    }

    return true;
}

bool SGCONF::Connect()
{
    try {
        ip::tcp::resolver resolver(io_service);
        ip::tcp::resolver::query query(host, port);
        ip::tcp::resolver::iterator endpoint_iterator = resolver.resolve(query);
        ip::tcp::resolver::iterator end;
        boost::system::error_code error = error::host_not_found;

        while (error && endpoint_iterator != end) {
            std::cerr << "Trying " << endpoint_iterator->endpoint().address().to_string() << ":" << port << "\n";
            socket.close();
            socket.connect(*endpoint_iterator++, error);
        }

        if (error) {
            throw boost::system::system_error(error);
        }

    }
    catch (std::exception & ex) {
        std::cerr << "Exception: " << ex.what() << "\n";
        return false;
    }

    return true;
}

void SGCONF::HandleHeader(const boost::system::error_code & error, size_t received)
{
    timer.cancel();
    if (error) {
        throw boost::system::system_error(error);
    }

    if (received < sizeof(respHeader)) {
        throw std::logic_error("Incomplete response header");
    }

    if (strncmp(respHeader.magic, PROTO_MAGIC, sizeof(respHeader.magic))) {
        throw std::logic_error("Invalid magic code in header");
    }

    uint32_t version = ntohl(respHeader.version);
    if (version > (2 << 8 | 0)) {
        throw std::logic_error("Unsupported protocol version");
    }

    switch (respHeader.code) {
        case RESP::INVALID_MAGIC:
            throw std::logic_error("Invalid magic");
        case RESP::UNSUPPORTED_VERSION:
            throw std::logic_error("Unsupported protocol version");
        case RESP::INVALID_CREDENTIALS:
            throw std::logic_error("Invalid login or password");
    };

    async_read(socket,
               boost::asio::buffer(cryptoRespBuffer, sizeof(cryptoRespBuffer)),
               boost::bind(&SGCONF::HandleCryptoHeader,
                           this,
                           boost::asio::placeholders::error,
                           boost::asio::placeholders::bytes_transferred));
    timer.expires_from_now(boost::posix_time::seconds(5));
    timer.async_wait(boost::bind(&SGCONF::HandleHeadersTimeout, this, boost::asio::placeholders::error));
}

void SGCONF::HandleCryptoHeader(const boost::system::error_code & error, size_t received)
{
    timer.cancel();
    if (error) {
        throw boost::system::system_error(error);
    }

    if (received < sizeof(respCryptoHeader)) {
        throw std::logic_error("Incomplete crypto header");
    }

    BF_cbc_encrypt(cryptoRespBuffer, reinterpret_cast<unsigned char *>(&respCryptoHeader), sizeof(respCryptoHeader), &key, iv.get(), BF_DECRYPT);

    respCryptoHeader.login[sizeof(respCryptoHeader.login) - 1] = 0;

    dataSize = ntohl(respCryptoHeader.dataSize);

    std::string cryptoLogin(respCryptoHeader.login);

    if (login != cryptoLogin) {
        throw std::logic_error("Password is invalid");
    }

    if (dataSize % 8) {
        throw std::logic_error("Invalid data size");
    }

    size_t size = 4096;

    if (dataSize < 4096) {
        size = dataSize;
    }

    xml = "";
    dataReceived = 0;

    async_read(socket,
               boost::asio::buffer(cryptoXMLBlock.get(), size),
               boost::bind(&SGCONF::HandleXML,
                           this,
                           boost::asio::placeholders::error,
                           boost::asio::placeholders::bytes_transferred));
    timer.expires_from_now(boost::posix_time::seconds(5));
    timer.async_wait(boost::bind(&SGCONF::HandleXMLTimeout, this, boost::asio::placeholders::error));
}

void SGCONF::HandleHeadersTimeout(const boost::system::error_code & error)
{
    if (error) {
        if (error != error::operation_aborted) {
            throw boost::system::system_error(error);
        }
        return;
    }
    throw std::logic_error("No headers within 5 sec");
}

void SGCONF::HandleXML(const boost::system::error_code & error, size_t received)
{
    timer.cancel();
    if (error) {
        if (error != boost::asio::error::eof) {
            throw boost::system::system_error(error);
        }
    }

    BF_cbc_encrypt(cryptoXMLBlock.get(), XMLBlock.get(), received, &key, iv.get(), BF_DECRYPT);
    xml.append(reinterpret_cast<char *>(XMLBlock.get()), received);
    dataReceived += received;

    if (dataReceived < dataSize) {
        size_t size = 4096;
        if (dataSize - dataReceived < 4096) {
            size = dataSize - dataReceived;
        }
        async_read(socket,
                   boost::asio::buffer(cryptoXMLBlock.get(), size),
                   boost::bind(&SGCONF::HandleXML,
                               this,
                               boost::asio::placeholders::error,
                               boost::asio::placeholders::bytes_transferred));
        timer.expires_from_now(boost::posix_time::seconds(5));
        timer.async_wait(boost::bind(&SGCONF::HandleXMLTimeout, this, boost::asio::placeholders::error));
    }
}

void SGCONF::HandleXMLTimeout(const boost::system::error_code& error)
{
    if (error) {
        if (error != error::operation_aborted) {
            throw boost::system::system_error(error);
        }
        return;
    }
    throw std::logic_error("No XML data within 5 sec");
}
